import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
import numpy as np

def getTrainSetAndTestSet(DataPath):
    data = pd.read_csv(DataPath)  # 读取CSV文件
    X = data[['AT', 'V', 'AP', 'RH']]  # 选取AT, V, AP, RH作为特征
    y = data[['PE']]  # 选取PE作为标签
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=1)  # 随机划分训练集和测试集，默认25%作为测试集
    print("训练集特征维度:", X_train.shape)
    print("训练集标签维度:", y_train.shape)
    print("测试集特征维度:", X_test.shape)
    print("测试集标签维度:", y_test.shape)
    return X_train, X_test, y_train, y_test

def TrainLinearRegression(X_train, y_train):
    linreg = LinearRegression()  # 创建线性回归模型
    linreg.fit(X_train, y_train)  # 使用训练集数据训练模型
    print("线性回归截距:", linreg.intercept_)
    print("线性回归系数:", linreg.coef_)
    return linreg

def EvaluationModel(linreg, X_test, y_test):
    y_pred = linreg.predict(X_test)  # 使用模型预测测试集
    mse = np.mean((y_pred - y_test) ** 2)  # 计算均方误差
    print("均方误差MSE:", mse)
    rmse = np.sqrt(mse)  # 计算均方根误差
    print("均方根误差RMSE:", rmse)
    return y_pred

def Visualization(y_test, y_pred):
    fig, ax = plt.subplots()  # 创建图形和轴
    ax.scatter(y_test, y_pred)  # 绘制散点图
    ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=5)  # 绘制对角线
    ax.set_xlabel("Measured")  # 设置x轴标签
    ax.set_ylabel("Predicted")  # 设置y轴标签
    display(fig)  # 使用display函数确保图形被正确显示
    plt.close(fig)  # 关闭图形以避免重复显示

if __name__ == "__main__":
    data_path = "Folds5x2_pp.csv"  # 数据集路径
    X_train, X_test, y_train, y_test = getTrainSetAndTestSet(data_path)  # 获取训练集和测试集
    linreg_model = TrainLinearRegression(X_train, y_train)  # 训练线性回归模型
    y_pred = EvaluationModel(linreg_model, X_test, y_test)  # 评估模型性能
    Visualization(y_test, y_pred)  # 可视化结果